/* Copyright (c) 2011 Alexandre Défossez
 * This file is part of HINAPT.
 *
 * HINAPT is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * HINAPT is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with HINAPT.  If not, see <http://www.gnu.org/licenses/>.
 */


#include "AffineTracker.hpp"

namespace {
  /*!
    Taile du bord interdit de l'image, en pixels
  */
  const int BORDER = 2;
}

const double AffineTracker::TRACK_FAILED = -1.;

AffineTracker::AffineTracker(int n, double threshold):
  n(n), threshold(threshold)
{}

void AffineTracker::getOriginal(Mat &m){
  m = original;
}

Point2f AffineTracker::getPoint(){
  return p;
}

void AffineTracker::update(Mat &img, Point2f point){
  original = img;
  p = point;

  intensity = norm(original, NORM_L1);
  intensity *= intensity;
}

void AffineTracker::calcA(Mat &img, Mat &grad, Mat &a){
  a.create(6, 1, CV_64FC1);
  a = Scalar(0);
  for(int x=BORDER; x<img.cols-BORDER; x++){
    for(int y=BORDER; y<img.rows-BORDER; y++){
      double d = static_cast<double>(original.at<uchar>(y,x)) - img.at<double>(y,x);
      double gx = grad.at<Vec2d>(y,x)[0];
      double gy = grad.at<Vec2d>(y,x)[1];
      double px = x-img.cols/2;
      double py =  y-img.rows/2;
      a.at<double>(0, 0) += d*px*gx;
      a.at<double>(1, 0) += d*px*gy;
      a.at<double>(2, 0) += d*py*gx;
      a.at<double>(3, 0) += d*py*gy;
      a.at<double>(4, 0) += d*gx;
      a.at<double>(5, 0) += d*gy;
    }
  }

}

void AffineTracker::calcT(Mat &grad, Mat &T){
  Size size = grad.size();
  T.create(6, 6, CV_64FC1);
  T = Scalar(0);
  for(int x=BORDER; x<size.width-BORDER; x++){
    for(int y=BORDER; y<size.height-BORDER; y++){
      double gx = grad.at<Vec2d>(y,x)[0];
      double gy = grad.at<Vec2d>(y,x)[1];
      double px = x-size.width/2;
      double py = y-size.height/2;
      T.at<double>(0, 0) += px*px*gx*gx;
      T.at<double>(0, 1) += px*px*gx*gy;
      T.at<double>(0, 2) += px*py*gx*gx;
      T.at<double>(0, 3) += px*py*gx*gy;
      T.at<double>(0, 4) += px*gx*gx;
      T.at<double>(0, 5) += px*gx*gy;

      T.at<double>(1, 1) += px*px*gy*gy;
      T.at<double>(1, 2) += px*py*gx*gy;
      T.at<double>(1, 3) += px*py*gy*gy;
      T.at<double>(1, 4) += px*gx*gy;
      T.at<double>(1, 5) += px*gy*gy;

      T.at<double>(2, 2) += py*py*gx*gx;
      T.at<double>(2, 3) += py*py*gx*gy;
      T.at<double>(2, 4) += py*gx*gx;
      T.at<double>(2, 5) += py*gx*gy;

      T.at<double>(3, 3) += py*py*gy*gy;
      T.at<double>(3, 4) += py*gx*gy;
      T.at<double>(3, 5) += py*gy*gy;

      T.at<double>(4, 4) += gx*gx;
      T.at<double>(4, 5) += gx*gy;
      T.at<double>(5, 5) += gy*gy;
    }
  }
  for(int y=0; y<6; y++){
    for(int x=0; x<y; x++){
      T.at<double>(y, x) = T.at<double>(x, y);
    }
  }
}

double AffineTracker::calcErr(Mat &window){
  double err = 0;
  for(int x=0; x<window.cols; x++){
    for(int y=0; y<window.rows; y++){
      double d=window.at<double>(y,x) - static_cast<double>(original.at<uchar>(y, x));
      err += d*d;
    }
  }
  return err/intensity;
}

bool AffineTracker::extractWindow(Mat &img, Mat &window, Mat &A, Mat &d){
  Rect boundaries(1, 1, img.cols-1, img.rows-1);
  Size size = window.size();
  Point2f point;
  for(int x=0; x<size.width; x++){
    for(int y=0; y<size.height; y++){
      point = applyAffineTransform(Point2f(x-size.width/2,
                                            y-size.height/2), A, d);
      point.x += p.x;
      point.y += p.y;
      if (!boundaries.contains(point)){
        return false;
      }
      window.at<double>(y,x) = getBrightness<uchar>(img, point.x, point.y);
    }
  }
  return true;
}

bool AffineTracker::convergence(Point2f lastWindowPos[],
                                Point2f windowPos[]){
  for(int i=0; i<4; i++){
    if (distance<Point2f>(lastWindowPos[i], windowPos[i])>threshold)
      return false;
  }
  return true;
}

void AffineTracker::windowMovement(Mat &A, Mat &d, Point2f pts[]){
  Size size = original.size();
  pts[0] = applyAffineTransform(Point2f(-size.width/2, -size.height/2), A, d);
  pts[1] = applyAffineTransform(Point2f(size.width/2, -size.height/2), A, d);
  pts[2] = applyAffineTransform(Point2f(size.width/2, size.height/2), A, d);
  pts[3] = applyAffineTransform(Point2f(-size.width/2, size.height/2), A, d);
}

double AffineTracker::internalTrack(Mat &img, Mat &A, Mat &d){
  Mat z(6, 1, CV_64FC1);
  Mat a, T, grad;
  Mat tD(A.size(), A.type());
  Mat td(d.size(), d.type());
  Mat window(original.size(), CV_64FC1);
  Mat tmp;
  Point2f lastWindowPos[4], windowPos[4];
  double err=0;

  if(!extractWindow(img, window, A, d)){
    return TRACK_FAILED;
  }


  windowMovement(A, d, lastWindowPos);

  for(int i=0;i<n; i++){
    calcGradient<double>(window, grad);
    calcT(grad, T);

    calcA(window, grad, a);
    solve(T, a, z);

    tD.at<double>(0,0) = z.at<double>(0, 0);
    tD.at<double>(1,0) = z.at<double>(1, 0);
    tD.at<double>(0,1) = z.at<double>(2, 0);
    tD.at<double>(1,1) = z.at<double>(3, 0);
    td.at<double>(0,0) = z.at<double>(4, 0);
    td.at<double>(1,0) = z.at<double>(5, 0);
    d += A*td;
    tmp = A*tD;
    A += tmp;

    if (!extractWindow(img, window, A, d)){ //On est sorti de l'image
      err = TRACK_FAILED;
      break;
    }
    windowMovement(A, d, windowPos);
    if (convergence(lastWindowPos, windowPos)){
      break;
    }
    for (int j=0; j<4; j++)
      lastWindowPos[j] = windowPos[j];
  }
  if (err != TRACK_FAILED)
    err = calcErr(window);

  return err;
}

double AffineTracker::track(Mat &img, Mat &D, Vec2d &d,
                            Point2f hint){
  Mat A = D;
  Mat Md(2, 1, CV_64FC1);
  Md.at<double>(0, 0) = hint.x - p.x;
  Md.at<double>(1, 0) = hint.y - p.y;

  double err = internalTrack(img, A, Md);
  d[0] = Md.at<double>(0, 0);
  d[1] = Md.at<double>(1, 0);
  D = A - Mat::eye(2,2, CV_64FC1);
  return err;
}

double AffineTracker::error(Mat &img, Point2f hint){
  Mat D = Mat::eye(2, 2, CV_64FC1);
  Vec2d d;
  return track(img, D, d, hint);
}
